



library(mvtnorm)
library(arm)
library(BRugs)
library(R2OpenBUGS)
library(coda)
library(car)
library(foreign)


setwd("C:/Users/Kevin/Dropbox/iREDS/Data/Analysis/scales/S1/waic")


#### ****** IMPORTANT NOTE:
# To run this, you need to set eqn to the scale number you are evaluating.  
eqn <- 1



# Throughout, a 0 indicates the restricted model and a 1 indicates the unrestricted

# delta.theta is the outcome matrix (each column is outcome replicates from an individual)
coda.data<-read.openbugs(stem="delta.theta.1_")
results.dtheta.1<-as.data.frame(rbind(coda.data[[1]], coda.data[[2]], coda.data[[3]]))
coda.data<-read.openbugs(stem="delta.theta.0_")
results.dtheta.0<-as.data.frame(rbind(coda.data[[1]], coda.data[[2]], coda.data[[3]]))

# beta is the matrix of structural parameters -- we only use S?.alpha1 for the WAIC
coda.data<-read.openbugs(stem="beta.1_")
results.1.beta<-as.data.frame(rbind(coda.data[[1]], coda.data[[2]], coda.data[[3]]))
coda.data<-read.openbugs(stem="beta.0_")
results.0.beta<-as.data.frame(rbind(coda.data[[1]], coda.data[[2]], coda.data[[3]]))

# e is the random effect matrix from the unrestricted model (each column is e replicates from an individual)
coda.data<-read.openbugs(stem="e_")  
results.e<-as.data.frame(rbind(coda.data[[1]], coda.data[[2]], coda.data[[3]]))




# ************ BEGIN WAIC TOOL (Unrestricted) ***************************************************
# Code taken from http://kylehardman.com/BlogPosts/View/7
# y_i is one observation from y.
# beta is the whole matrix of samples from the posterior distribution.
likelihood.1 = function(y.i, beta, z, e, eqn) {   # x1 is resorg, x2 is the distance 
  #Get the individual parameter vectors out of beta.
  if (eqn == 1) {
    alpha1 = beta[,"S1.alpha1"]
	} else if (eqn == 2) {
    alpha1 = beta[,"S2.alpha1"]
	} else if (eqn == 3) {
    alpha1 = beta[,"S3.alpha1"]
	} else if (eqn == 5) {
    alpha1 = beta[,"S5.alpha1"]
	} else if (eqn == 6) {
    alpha1 = beta[,"S6.alpha1"]
	} else if (eqn == 8) {
    alpha1 = beta[,"S8.alpha1"]
	} else print("Error!")
  theta.bar <- (alpha1*z + e) 
  dnorm(y.i, theta.bar)
}
# y is the vector of data (this function would need to be modified slightly if y is a matrix)
# beta is a matrix of samples from the posterior distribution of the parameters.
# likeFun is a function that takes two arguments: 
#   a single observation from y and a matrix of posterior chains, beta. 
#   likeFun returns a vector of likelihoods, each corresponding to a row of beta.
#   For a function matching the requirements of likeFun, see likelihood_singleY, defined above.

calculate_WAIC.1 = function(y, beta, x1, x2, eqn, likeFun) { # here, x1=condition, x2==e
  
  # Initialize accumulator variables that will be added to
  LPPD = 0
  P = 0
  
  for (i in 1:length(y[1,])) {
    
    # Get a vector of likelihoods for this observation in y
    L = likeFun(y[,i], beta, x1[i], x2[,i], eqn)  # the respondent indexes yield the right column
    #print(L)
    # Calculate LPPD for this observation
    LPPD = LPPD + log( mean(L) )
    
    # Calculate P
    P = P + var( log(L) )
    #print(P)
    
  }
  
  
  # Calculate the WAIC
  WAIC = -2 * (LPPD - P)
  
  # Return a list containing the results
  list(WAIC = WAIC, P = P, LPPD = LPPD)
}


waic.1 = calculate_WAIC.1(results.dtheta.1, results.1.beta, dataset$condition, results.e, eqn, likelihood.1) #eqn is set to scale number






# ************ BEGIN WAIC TOOL (restricted) ***************************************************
# Code taken from http://kylehardman.com/BlogPosts/View/7
# y_i is one observation from y.
# beta is the whole matrix of samples from the posterior distribution.
likelihood.0 = function(y.i, beta, z, eqn) {   # x1 is resorg, x2 is the distance 
  #Get the individual parameter vectors out of beta.
  if (eqn == 1) {
    alpha1 = beta[,"S1.alpha1"]
  } else if (eqn == 2) {
    alpha1 = beta[,"S2.alpha1"]
  } else if (eqn == 3) {
    alpha1 = beta[,"S3.alpha1"]
  } else if (eqn == 5) {
    alpha1 = beta[,"S5.alpha1"]
  } else if (eqn == 6) {
    alpha1 = beta[,"S6.alpha1"]
  } else if (eqn == 8) {
    alpha1 = beta[,"S8.alpha1"]
  } else print("Error!")
  theta.bar <- (alpha1*z) 
  dnorm(y.i, theta.bar)
}
# y is the vector of data (this function would need to be modified slightly if y is a matrix)
# beta is a matrix of samples from the posterior distribution of the parameters.
# likeFun is a function that takes two arguments: 
#   a single observation from y and a matrix of posterior chains, beta. 
#   likeFun returns a vector of likelihoods, each corresponding to a row of beta.
#   For a function matching the requirements of likeFun, see likelihood_singleY, defined above.

calculate_WAIC.0 = function(y, beta, x1, eqn, likeFun) { # here, x1=condition
  
  # Initialize accumulator variables that will be added to
  LPPD = 0
  P = 0
  
  for (i in 1:length(y[1,])) {
    
    # Get a vector of likelihoods for this observation in y
    L = likeFun(y[,i], beta, x1[i], eqn)  # the respondent indexes yield the right column
    #print(L)
    # Calculate LPPD for this observation
    LPPD = LPPD + log( mean(L) )
    
    # Calculate P
    P = P + var( log(L) )
    #print(P)
    
  }
  
  
  # Calculate the WAIC
  WAIC = -2 * (LPPD - P)
  
  # Return a list containing the results
  list(WAIC = WAIC, P = P, LPPD = LPPD)
}


waic.0 = calculate_WAIC.0(results.dtheta.0, results.0.beta, dataset$condition, eqn, likelihood.0) #eqn is set to scale number




sink("WAIC.txt")
print(waic.1)
print(waic.0)
sink()



